package com.ds.api;

import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;

/**
 * @Author 21326
 * @Date 2025 2025/11/9 21:25
 */
@RestController
public class LoginController {

    // 模拟用户数据库，实际应从数据库或配置中获取
    private static final Map<String, String> USER_PASSWORD_MAP = new HashMap<>();
    private static final String QOP = "auth";

    static {
        USER_PASSWORD_MAP.put("admin", "123456");
    }

    public enum Algorithm {
        MD5("MD5", "MD5"),
        MD5_SESS("MD5-sess", "MD5-sess"),
        SHA_256("SHA-256", "SHA-256"),
        SHA_256_SESS("SHA-256-sess", "SHA-256-sess"),
        SHA_512_256("SHA-512-256", "SHA-512-256"),
        SHA_512_256_SESS("SHA-512-256-sess", "SHA-512-256-sess");

        private final String desc;
        private final String value;

        Algorithm(String desc, String value) {
            this.desc = desc;
            this.value = value;
        }

        public String getDesc() {
            return desc;
        }

        public String getValue() {
            return value;
        }

        public static Algorithm fromValue(String value) {
            for (Algorithm alg : values()) {
                if (alg.value.equalsIgnoreCase(value)) {
                    return alg;
                }
            }
            return MD5; // 默认使用 MD5
        }
    }

    @PostMapping("/login")
    public ObjectNode login(HttpServletRequest request, HttpServletResponse response) {
        ObjectNode node = new ObjectNode(JsonNodeFactory.instance);
        String authorization = request.getHeader("authorization");

        if (StringUtils.hasText(authorization) && authorization.startsWith("Digest ")) {
            try {
                System.out.println("Authorization +> " + authorization);
                Map<String, String> authParams = parseDigestAuthorization(authorization.substring(7));
                String username = authParams.get("username");
                String nonce = authParams.get("nonce");
                String uri = authParams.get("uri");
                String responseDigest = authParams.get("response");
                String qop = authParams.get("qop");
                String nc = authParams.get("nc");
                String cnonce = authParams.get("cnonce");
                String algorithmStr = authParams.getOrDefault("algorithm", "MD5");
                Algorithm algorithm = Algorithm.fromValue(algorithmStr);

                // 1. 验证用户是否存在
                String password = USER_PASSWORD_MAP.get(username);
                if (Objects.isNull(password)) {
                    sendUnauthorized(response, node, algorithm);
                    return node;
                }

                // 2. 计算 HA1
                String ha1 = computeHa1(username, password, nonce, cnonce, algorithm);
                System.out.println("HA1 +> " + ha1);

                // 3. 计算 HA2
                String ha2 = computeHa2(request.getMethod(), uri, algorithm);
                System.out.println("HA2 +> " + ha2);

                // 4. 计算最终摘要
                String ha3 = computeResponse(ha1, nonce, nc, cnonce, qop, ha2, algorithm);
                System.out.println("HA3 +> " + ha3);

                // 5. 比较摘要
                System.out.println("Response Digest +> " + responseDigest);
                if (ha3.equals(responseDigest)) {
                    node.put("code", 200);
                    node.put("msg", "AUTH SUCCESS");
                    node.put("data", "{\"token\":\"JWT Token\"}");
                } else {
                    sendUnauthorized(response, node, algorithm);
                }
            } catch (Exception e) {
                sendUnauthorized(response, node, Algorithm.MD5);
            }
        } else {
            sendUnauthorized(response, node, Algorithm.MD5);
        }

        return node;
    }

    private String computeHa1(String username, String password, String nonce, String cnonce, Algorithm algorithm) {
        String realm = "http://realm.com.cn";
        String str1 = username + ":" + realm + ":" + password;
        String ha1 = digest(str1, algorithm.value);

        if (algorithm == Algorithm.MD5_SESS ||
                algorithm == Algorithm.SHA_256_SESS ||
                algorithm == Algorithm.SHA_512_256_SESS) {
            String str2 = ha1 + ":" + nonce + ":" + cnonce;
            ha1 = digest(str2, algorithm.value);
        }

        return ha1;
    }

    private String computeHa2(String method, String uri, Algorithm algorithm) {
        String str2 = method + ":" + uri;
        return digest(str2, algorithm.value);
    }

    private String computeResponse(String ha1, String nonce, String nc, String cnonce, String qop, String ha2, Algorithm algorithm) {
        String str3;
        if ("auth-int".equals(qop)) {
            str3 = ha1 + ":" + nonce + ":" + nc + ":" + cnonce + ":" + qop + ":" + ha2;
        } else if ("auth".equals(qop)) {
            str3 = ha1 + ":" + nonce + ":" + ha2;
        } else {
            str3 = ha1 + ":" + nonce + ":" + ha2;
        }
        return digest(str3, algorithm.value);
    }

    private String digest(String data, String algorithm) {
        try {
            String algName = algorithm.toUpperCase();
            MessageDigest md;

            switch (algName) {
                case "MD5":
                    md = MessageDigest.getInstance("MD5");
                    break;
                case "SHA-256":
                    md = MessageDigest.getInstance("SHA-256");
                    break;
                case "SHA-512-256":
                    md = MessageDigest.getInstance("SHA-512/256");
                    break;
                default:
                    throw new IllegalArgumentException("Unsupported algorithm: " + algorithm);
            }

            byte[] digest = md.digest(data.getBytes(StandardCharsets.UTF_8));
            StringBuilder sb = new StringBuilder();
            for (byte b : digest) {
                sb.append(String.format("%02x", b));
            }
            return sb.toString();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void sendUnauthorized(HttpServletResponse response, ObjectNode node, Algorithm algorithm) {
        node.put("code", 401);
        node.put("msg", "NOT AUTH");
        node.put("data", "[]");
        String nonce = UUID.randomUUID().toString().replace("-", "");
        String www_authorization = String.format("Digest realm=\"%s\", qop=\"%s\", nonce=\"%s\", algorithm=\"%s\"",
                "http://realm.com.cn", QOP, nonce, algorithm.getValue());
        response.setHeader("WWW-Authenticate", www_authorization);
        response.setStatus(401);
    }

    private Map<String, String> parseDigestAuthorization(String authHeader) {
        Map<String, String> params = new HashMap<>();
        String[] tokens = authHeader.split(",(?=(?:[^\"]*\"[^\"]*\")*[^\"]*$)");
        for (String token : tokens) {
            String[] keyValue = token.split("=", 2);
            if (keyValue.length == 2) {
                String key = keyValue[0].trim();
                String value = keyValue[1].trim().replace("\"", "");
                params.put(key, value);
            }
        }
        return params;
    }
}